# -*- coding: utf-8 -*-
import os
from types import SimpleNamespace

import cv2
import torch
import numpy as np
import matplotlib.cm as cm
from kornia.geometry.transform import warp_perspective
import os.path as osp
from load_model import load_model
from src.utils.plotting import make_matching_figure


def save_matching_figure(path, img0, img1, mkpts0, mkpts1, inlier_mask, color):
    """保存RANSAC过滤后的匹配结果图"""
    if inlier_mask is None or len(inlier_mask) == 0:
        print("⚠️ No inlier mask found, skip saving.")
        return
    inlier_mask = inlier_mask.astype(bool).squeeze()
    mkpts0_inliers = mkpts0[inlier_mask]
    mkpts1_inliers = mkpts1[inlier_mask]
    color = color[inlier_mask]
    text = [f'Matches: {len(mkpts0_inliers)}']
    make_matching_figure(img0, img1, mkpts0_inliers, mkpts1_inliers,
                         color, text=text, path=path, dpi=150)
    print(f"✅ Saved figure: {path}")


def save_matching_figure2(path, img0, img1, mkpts0, mkpts1, color):
    """保存RANSAC过滤前的匹配结果图"""
    text = [f'Matches: {len(mkpts0)}']
    make_matching_figure(img0, img1, mkpts0, mkpts1,
                         color, text=text, path=path, dpi=150)
    print(f"✅ Saved figure: {path}")


def H_transform(img_tensor, homography):
    """使用Kornia透视变换对齐图像"""
    shape = img_tensor.shape[2:]
    warped = warp_perspective(img_tensor, homography, shape, align_corners=True)
    return warped


def evaluate_pair(method, img_path1, img_path2, save_dir, save_figs=True):
    """主流程：加载模型→匹配→估计单应性→保存结果"""

    # 1. 加载匹配模型
    print(f"\n🔧 Using method: {method}")

    # ============================
    # 参数配置表
    # ============================
    method_args = {
        "xoftr": {
            "ckpt": "./weights/minima_xoftr.ckpt",
            "match_threshold": 0.3,
            "fine_threshold": 0.1,
        },
        "loftr": {
            "ckpt": "./weights/minima_loftr.ckpt",
            "thr": 0.2,
        },
        "sp_lg": {
            "ckpt": "./weights/minima_lightglue.pth",
        },
        "roma": {
            "ckpt": "./weights/minima_roma.pth",
            "ckpt2": "tiny",#large，可以不写这个残念书然后看报错
        },
    }

    if method not in method_args:
        raise ValueError(f"❌ Unsupported method: {method}")

    # 转成命名空间对象传给 load_model
    args = SimpleNamespace(**method_args[method])

    # ============================
    # 加载匹配器
    # ============================
    matcher = load_model(method, args)
    # 2. 读取图片
    img0 = cv2.imread(img_path1)
    img1 = cv2.imread(img_path2)
    img0_rgb = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB)
    img1_rgb = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)

    # 3. 执行匹配
    print("🔍 Running feature matching ...")
    match_res = matcher(img_path1, img_path2)

    mkpts0 = match_res['mkpts0']#图像1中匹配点的坐标 (x, y)
    mkpts1 = match_res['mkpts1']#图像2中对应匹配点的坐标 (x, y)
    mconf = match_res['mconf']# # 每个匹配点的置信度（匹配可信度）
    #图像1中坐标为 (x₁, y₁) 的点，与图像2中 (x₂, y₂) 的点是同一个真实位置（例如建筑的角、物体的特征点等）。

    # 输出特征点信息
    print(f"\n📷 图像1特征点数量: {len(mkpts0)}")
    print(f"📷 图像2特征点数量: {len(mkpts1)}")
    if len(mkpts0) > 0:
        print(f"示例匹配点 (前3个):")
        for i in range(min(3, len(mkpts0))):
            print(f"  P{i+1}: img1({mkpts0[i]}) -> img2({mkpts1[i]})")

    # 归一化置信度并转为颜色
    if len(mconf) > 0:
        mconf = (mconf - mconf.min()) / (mconf.max() - mconf.min() + 1e-5)
    color = cm.jet(mconf)

    # 4. RANSAC计算单应矩阵
    #RANSAC（随机采样一致性算法）是一个过滤错误匹配点的算法。因为特征匹配模型并不是完美的，它可能会误判某些点。
    """
    ✅ RANSAC 会做的事： 
    随机选几个匹配点、计算一个单应性矩阵 H（描述图像1 → 图像2 的透视变换关系）、用 H 去预测其他匹配点的位置
    、保留符合变换规律的点（内点 inliers）、 重复多次，找出最大一致集合。最终它输出：ret_H  # 3x3 单应性矩阵、inliers  # True/False mask, 表示哪些匹配是可靠的
    H 是一个描述两张平面图像之间透视变换的 3×3 矩阵，它能把第一张图的坐标点 (x, y) 映射到第二张图中的对应位置 (x’, y’)。
    这在 图像拼接、全景图、SLAM、相机姿态估计 等任务中都非常重要。
    """
    if len(mkpts0) >= 4:
        H, inliers = cv2.findHomography(mkpts0, mkpts1, cv2.RANSAC)
        num_inliers = inliers.sum() if inliers is not None else 0
        print(f"\n✅ RANSAC完成: {num_inliers} 个内点")
    else:
        print("⚠️ 匹配点少于4个，无法估计单应性矩阵")
        H, inliers = None, None

    # 5. 绘制并保存结果
    if save_figs:
        os.makedirs(save_dir, exist_ok=True)
        name0 = osp.splitext(osp.basename(img_path1))[0]
        name1 = osp.splitext(osp.basename(img_path2))[0]

        path_before = osp.join(save_dir, f"{name0}_{name1}_before_ransac_{method}.jpg")
        path_after = osp.join(save_dir, f"{name0}_{name1}_after_ransac_{method}.jpg")
        #所有模型匹配点（包括错误匹配）
        save_matching_figure2(path_before, img0_rgb, img1_rgb, mkpts0, mkpts1, color)
        if inliers is not None:
            ##保存RANSAC过滤后的正确匹配点
            save_matching_figure(path_after, img0_rgb, img1_rgb, mkpts0, mkpts1, inliers, color)

    # 6. 透视变换可视化
    #保存图像1透视变换后的结果（几何上对齐图像2）
    if H is not None:
        print("\n🌀 应用单应性矩阵变换图像1")
        im0_tensor = torch.tensor(img0_rgb, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) / 255.
        H_tensor = torch.tensor(H, dtype=torch.float32).unsqueeze(0)
        #把“图像1”通过刚刚算出的单应性矩阵 H 进行几何变换，让它尽量对齐到“图像2”的视角。
        warped_tensor = H_transform(im0_tensor, H_tensor)
        warped_img = (warped_tensor.squeeze().permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
        path_warped = osp.join(save_dir, f"{name0}_warped_{method}.jpg")
        cv2.imwrite(path_warped, cv2.cvtColor(warped_img, cv2.COLOR_RGB2BGR))
        print(f"✅ Saved warped image: {path_warped}")

if __name__ == '__main__':
    # ============================
    #roma：
    # C:\Users\ASUS/.cache\torch\hub\checkpoints\dinov2_vitl14_pretrain.pth
    #xoftr：
    #C:\Users\ASUS/.cache\torch\hub\checkpoints\superpoint_v1.pth
    # ============================
    method = "roma"  # 可选：'loftr', 'sp_lg', 'xoftr', 'roma'
    img_path1 = "./demo/vis_test.png"
    img_path2 = "./demo/depth_test.png"
    save_dir = "./demo/results"
    save_figs = True

    # ============================
    # 🧭 执行流程
    # ============================
    evaluate_pair(method, img_path1, img_path2, save_dir, save_figs)
